package top.wshape1.rpc.registry;

import io.netty.util.concurrent.DefaultThreadFactory;
import lombok.extern.slf4j.Slf4j;
import redis.clients.jedis.Jedis;
import redis.clients.jedis.JedisPool;
import redis.clients.jedis.JedisPubSub;
import top.wshape1.rpc.common.entity.ServiceMeta;
import top.wshape1.rpc.common.constant.RoleType;
import top.wshape1.rpc.config.JJRpcConfig;
import top.wshape1.rpc.config.RegistryConfig;

import java.util.*;
import java.util.concurrent.*;
import java.util.function.Function;

/**
 * @author Wshape1
 * @version 1.0
 * @description Redis注册中心
 * @since 2024-02-02
 */

@Slf4j
public class RedisRegistry implements RpcRegistry {

    /*
     *  Redis config
     */
    private final RegistryConfig.RedisRegistryConfig REDIS_CONFIG = JJRpcConfig.CONFIG.getRegistry().getRedis();
    private final String KEY_PREFIX = REDIS_CONFIG.getPrefix();

    /*
     *  Constants
     */

    private static final int DEFAULT_EXPIRE_PERIOD = 60 * 1000;
    private static final String REGISTER = "register";
    private static final String UNREGISTER = "unregister";

    /*
     *  Functional Value
     */

    /**
     * （消费者）缓存每个订阅的服务接口的生产者地址
     * serviceName -> Set<ServiceMeta>
     */
    private final ConcurrentHashMap<String, Set<ServiceMeta>> SERVICE_CACHE = new ConcurrentHashMap<>();

    /**
     * （消费者）存储每一个消费者对某个服务的订阅
     * serviceName -> Subscriber
     */
    private final ConcurrentHashMap<String, Subscriber> SUBSCRIBER_MAP = new ConcurrentHashMap<>();

    /**
     * （生产者）存储已经注册的服务
     *  value: serviceName
     */
    private final CopyOnWriteArraySet<String> REGISTERED_PROVIDER = new CopyOnWriteArraySet<>();

    private final ScheduledExecutorService expireExecutor = Executors.newScheduledThreadPool(1, new DefaultThreadFactory("RpcRegistryExpireTimer", true));

    private final ScheduledFuture<?> expireFuture;

    /*
     *  Jedis Client
     */

    private final JedisPool jedisPool;

    /*
     *  Constructor
     */

    public RedisRegistry() {
        // 创建Jedis连接对象
        jedisPool = new JedisPool(REDIS_CONFIG.getPool(),
                REDIS_CONFIG.getHost(), REDIS_CONFIG.getPort(),
                REDIS_CONFIG.getSoTimeout(), REDIS_CONFIG.getConnectionTimeout(),
                REDIS_CONFIG.getUsername(), REDIS_CONFIG.getPassword(),
                REDIS_CONFIG.getDb(), null,
                REDIS_CONFIG.isSsl(), null, null, null);

        // 定时 保活生产者服务
        this.expireFuture = expireExecutor.scheduleWithFixedDelay(() -> {
            try {
                execute(j -> {
                    ServiceMeta serviceMeta = ServiceMeta.selfServiceMeta(null, RoleType.provider);
                    String value = serviceMeta.getAddress();
                    // 对本机注册的生产者服务进行续期
                    REGISTERED_PROVIDER.forEach((serviceName) -> {
                        serviceMeta.setServiceName(serviceName);
                        String key = toCategoryPath(serviceMeta);

                        long now = System.currentTimeMillis();
                        String oldExpire = j.hget(key, value);
                        // 说明该服务没有过期, 可以续期
                        if (Long.parseLong(oldExpire) > now) {
                            String newExpire = String.valueOf(now + DEFAULT_EXPIRE_PERIOD);
                            j.hset(key, value, newExpire);
                            j.publish(key, REGISTER);
                        } else {
                            // 过期删除
                            j.hdel(key, value);
                        }
                    });

                    serviceMeta.setRole(RoleType.consumer);
                    // 对本机注册的消费者进行续期
                    SUBSCRIBER_MAP.forEach((serviceName, subscriber) -> {
                        serviceMeta.setServiceName(serviceName);
                        String key = toCategoryPath(serviceMeta);

                        long now = System.currentTimeMillis();
                        String oldExpire = j.hget(key, value);
                        // 说明该服务没有过期, 可以续期
                        if (Long.parseLong(oldExpire) > now) {
                            String newExpire = String.valueOf(now + DEFAULT_EXPIRE_PERIOD);
                            j.hset(key, value, newExpire);
                            j.publish(key, REGISTER);
                        } else {
                            // 过期删除
                            j.hdel(key, value);
                        }
                    });
                    return null;
                });
            } catch (Throwable t) {
                log.error("Unexpected exception occur at defer expire time, cause: " + t.getMessage(), t);
            }
        }, DEFAULT_EXPIRE_PERIOD / 2, DEFAULT_EXPIRE_PERIOD / 2, TimeUnit.MILLISECONDS);

    }

    /*
     *  Override Methods
     */

    @Override
    public void register(ServiceMeta serviceMeta) {
        String key = toCategoryPath(serviceMeta);
        String value = serviceMeta.getAddress();
        String expire = String.valueOf(System.currentTimeMillis() + DEFAULT_EXPIRE_PERIOD);
        try {
            execute(j -> {

                Long l = j.hset(key, value, expire);
                j.publish(key, REGISTER);
                return l;
            });
        } catch (Throwable t) {
            throw new RuntimeException("Failed to register service to redis registry. Address: "
                    + serviceMeta.getAddress()
                    + ", Service: " + serviceMeta.getServiceName()
                    + "Cause: " + t.getMessage(), t);
        }
        // 如果是消费者，自动订阅对应生产者服务
        if (serviceMeta.getRole() == RoleType.consumer)
            subscribe(serviceMeta.getServiceName());
        else
            REGISTERED_PROVIDER.add(serviceMeta.getServiceName());
    }

    @Override
    public void unregister(ServiceMeta serviceMeta) {
        String key = toCategoryPath(serviceMeta);
        String value = serviceMeta.getAddress();
        try {
            execute(j -> {

                Long l = j.hdel(key, value);
                j.publish(key, UNREGISTER);
                return l;

            });
        } catch (Throwable t) {
            throw new RuntimeException("Failed to unregister service to redis registry. Address: "
                    + serviceMeta.getAddress()
                    + ", Service: " + serviceMeta.getServiceName()
                    + "Cause: " + t.getMessage(), t);
        }
        if (serviceMeta.getRole() == RoleType.provider)
            REGISTERED_PROVIDER.remove(serviceMeta.getAddress());
        else
            unsubscribe(serviceMeta.getServiceName());
    }

    @Override
    public List<ServiceMeta> lookup(String serviceName) {
        // 每次生产者服务 续期、注册、注销都会发送事件，每个服务消费者接收到都会更新缓存
        // 可以保证缓存服务的可用性
        Set<ServiceMeta> set = SERVICE_CACHE.get(serviceName);
        if (set != null)
            return new ArrayList<>(set);
        return new ArrayList<>(0);
    }

    @Override
    public void subscribe(String serviceName) {
        String key = toCategoryPath4Sub(serviceName);

        // 订阅服务，不重复订阅
        Subscriber subscriber = SUBSCRIBER_MAP.get(serviceName);
        if (subscriber == null) {
            Subscriber newSubscriber = new Subscriber(serviceName);
            SUBSCRIBER_MAP.putIfAbsent(serviceName, newSubscriber);
            subscriber = SUBSCRIBER_MAP.get(serviceName);
            if (subscriber == newSubscriber) {
                subscriber.start();
            }
        }

        Map<String, String> map = execute(j -> {
            // 如果没有缓存过
            if (SERVICE_CACHE.get(serviceName) == null)
                return j.hgetAll(key);
            else
                return null;
        });
        // 把服务生产者都写入缓存
        if (map != null) {
            HashSet<ServiceMeta> set = toServiceMetaSet(serviceName, map);
            SERVICE_CACHE.putIfAbsent(serviceName, set);
        }
    }

    @Override
    public void unsubscribe(String serviceName) {
        Subscriber subscriber = SUBSCRIBER_MAP.remove(serviceName);
        if (subscriber != null) {
            subscriber.interrupt();
        }
    }

    @Override
    public void destroy() {
        try {
            expireFuture.cancel(true);
        } catch (Throwable t) {
            log.warn(t.getMessage(), t);
        }
        // 注销消费者
        ServiceMeta serviceMeta = ServiceMeta.selfServiceMeta(null, RoleType.consumer);
        try {
            SUBSCRIBER_MAP.forEach((serviceName, subscriber) -> {
                serviceMeta.setServiceName(serviceName);
                unregister(serviceMeta);
                unsubscribe(serviceName);
            });
        } catch (Throwable t) {
            log.warn(t.getMessage(), t);
        }
        // 注销生产者
        serviceMeta.setRole(RoleType.provider);
        try {
            REGISTERED_PROVIDER.forEach((serviceName) -> {
                serviceMeta.setServiceName(serviceName);
                unregister(serviceMeta);
            });
        } catch (Throwable t) {
            log.warn(t.getMessage(), t);
        }
        // close jedis pool
        try {
            if (!jedisPool.isClosed())
                jedisPool.close();
        } catch (Throwable t) {
            log.warn(t.getMessage(), t);
        }
        try {
            if (!expireExecutor.isShutdown())
                expireExecutor.shutdown();
        } catch (Throwable t) {
            log.warn(t.getMessage(), t);
        }
    }

    /*
     *  Private Methods
     */

    private String toCategoryPath4Sub(String serviceName) {
        return "/" + KEY_PREFIX + "/" + serviceName + "/" + RoleType.provider.name();
    }

    private String toCategoryPath(ServiceMeta serviceMeta) {
        return "/" + KEY_PREFIX + "/" + serviceMeta.getServiceName() + "/" + serviceMeta.getRole().name();
    }

    private Jedis getJedis() {
        return jedisPool.getResource();
    }

    private <R> R execute(Function<Jedis, R> function) {
        Jedis jedis = getJedis();
        R ret = function.apply(jedis);
        if (jedisPool != null)
            jedis.close();
        return ret;
    }

    private HashSet<ServiceMeta> toServiceMetaSet(String serviceName, Map<String, String> map) {
        HashSet<ServiceMeta> newSet = new HashSet<>();
        if (map != null && !map.isEmpty()) {
            map.forEach((k, v) -> {
                ServiceMeta meta = new ServiceMeta();
                meta.setServiceName(serviceName);
                meta.setAddress(k);
                meta.setRole(RoleType.provider);
                newSet.add(meta);
            });
        }
        return newSet;
    }

    /*
     *  Private Classes
     */
    private class Subscriber extends Thread {

        private final String serviceName;

        public Subscriber(String serviceName) {
            super.setDaemon(true);
            super.setName("RedisSubscribe-" + serviceName);
            this.serviceName = serviceName;
        }

        @Override
        public void run() {
            // 会阻塞，并一直占用此Jedis连接
            try {
                execute(j -> {
                    j.subscribe(new ServiceSub(), toCategoryPath4Sub(serviceName));
                    return null;
                });
            } catch (Throwable t) {
                log.warn("订阅服务{}关闭，原因：{}", serviceName, t.getMessage(), t);
            }
        }


    }

    private class ServiceSub extends JedisPubSub {

        @Override
        public void onMessage(String key, String message) {
            // 接收到register或unregister事件
            if (REGISTER.equals(message) || UNREGISTER.equals(message)) {
//                log.debug("接收到{}事件，更新生产者服务缓存", message);
                // 更新缓存列表
                SERVICE_CACHE.compute(key, (serviceName, set) -> {
                    Map<String, String> map = execute(j -> j.hgetAll(key));
                    return toServiceMetaSet(serviceName, map);
                });
            }
        }

        @Override
        public void onPMessage(String pattern, String channel, String message) {
            onMessage(channel, message);
        }
    }


}
